#!/usr/bin/ruby
# Generates a template for exploit_config.h
# This script is heaviliy adapted for VHBL: it assumes that perfect syscalls are not available, and that hbl needs to be run in a "flat" folder
#
# Required input files are:
#
# 1) "memdump.bin", a user memory dump from PSPLink with 
# "savemem 0x08800000 0x01800000 memdump.bin"
#
# 2) from the same game session as the memdump: "uidlist.txt", which is the the output of the "uidlist"
# command in PSPLink
#
# 3) "lwmtxlist.txt", which is the output of the "lwmtxlist" command in PSPLink Mod by 173210
# You don't need it if the game doesn't have lwmutex.
#
# Initial code by JJS & wololo

class Integer
	def to_a
		[self]
	end
end

def help
    puts 'Usage: ' + $0 + %Q[ <OUTPUT>

You need to provide 3 files: memdump.bin, uidlist.txt, and lwmtxlist.txt
  - memdump.bin: a user memory dump from PSPLink
       ("savemem 0x08800000 0x01800000 memdump.bin")
  - from the same game session as the memdump: "uidlist.txt",
       which is the the output of the "uidlist" command in PSPLink
  - "lwmtxlist.txt":
       the output of the "lwmtxlist" command in PSPLink Mod by 173210.
       You don't need it if the game doesn't have lwmutex.
]
end
config_keys = [
	'TH_ADDR_LIST',
	'ALARM_ADDR_LIST',
	'LWMUTEX_ADDR_LIST',
	'EV_ADDR_LIST',
	'SEMA_ADDR_LIST',
	'VPL_ADDR_LIST',
	'FPL_ADDR_LIST',
	'GAME_FREEMEM_ADDR'
]

config_values = { }

config_keys.each { |key|
	config_values.store(key, [])
}

if (!ARGV[0])
	help()
	abort()
end

if (File.exists?('memdump.bin') && File.exists?('uidlist.txt'))
	#
	# Find Semaphore / Thread / evlist / game addresses
	#
	base = 0x08800000
	ent_type = nil
	lwmtx = false

	memdump = IO.binread('memdump.bin')
	IO.foreach('uidlist.txt') { |line|
		#[Fpl]    UID 0x00288E0D (attr 0x0 entry 0x88014470)
		if line =~ /^\[(.+)\].+$/
			ent_type = $1
			next
		end

		#(UID): 0x048E3761, (entry): 0x882471b8, (size): 48, (attr): 0xFF, (Name): sgx-ao-evf
		line.scan(/\(UID\): (0x.+?),.*\(attr\): (0x.+?),.*\(Name\): (.+)$/) {
			|ent_uid, ent_attr, ent_name|
			# skip kernel objects
			if ent_attr != '0xFF'
				break
			end

			index = memdump.index(ent_uid.hex.to_a.pack('V'))
			if !index
				puts "can't find address for UID " + ent_uid;
				break
			end

			value = "0x0#{(index + base).to_s(16).upcase}"
                
			if ent_type == 'Thread'
				# skip main thread
				if !ent_name.include?('main')
					config_values['TH_ADDR_LIST'] << value
				end
			elsif ent_type == 'Alarm'
				config_values['ALARM_ADDR_LIST'] << value
			elsif ent_type == 'LwMutex'
				lwmtx = true
			elsif ent_type == 'EventFlag'
				config_values['EV_ADDR_LIST'] << value
			elsif ent_type == 'Semaphore'
				config_values['SEMA_ADDR_LIST'] << value  
			elsif ent_type == 'SceSysMemMemoryBlock'
				# skip thread stacks
				if !ent_name.include?('stack:')
					config_values['GAME_FREEMEM_ADDR'] << value
				end
			elsif ent_type == 'Vpl'
				config_values['VPL_ADDR_LIST'] << value  
			elsif ent_type == 'Fpl'
				config_values['FPL_ADDR_LIST'] << value  
			end
		}
	}

    if (lwmtx)
            lwmtxlist = File.new('lwmtxlist.txt', 'r')

            #UID: 0xXXXXXXXX - Workarea: 0xYYYYYYYY - Name: ZZZZZZZZ
            regexpEntry = /^UID: .+ - Workarea: (0x.+) - Name: .+$/

            while (line = lwmtxlist.gets)
                line.chomp!
                if (line.match(regexpEntry)) 
                    line.scan(regexpEntry) { |addr|
                         config_values['LWMUTEX_ADDR_LIST'].push(addr)
                    }
                end
            end
    end
else
	help()
	abort()
end

if (config_values['TH_ADDR_LIST'].empty? &&
	config_values['ALARM_ADDR_LIST'].empty? &&
	config_values['LWMUTEX_ADDR_LIST'].empty? &&
	config_values['EV_ADDR_LIST'].empty? &&
	config_values['SEMA_ADDR_LIST'].empty? &&
	config_values['GAME_FREEMEM_ADDR'].empty?)
	puts '!!!! ERROR, NO DATA FOUND! MAKE SURE THAT uidlist.txt IS IN UNIX FORMAT!!!'
	abort()
end

out = File.new(ARGV[0], 'w')

out.print %Q[// This is a pregenerated file made by gen_exploit_config.rb. HBL_ROOT is labelled as "TODO", you need to enter your path here

// Replace the "TODO" below with the address which the exploit loads binary to
#define LOADER_ADDR TODO
// Replace the "TODO" below with the code of your game (e.g UCUS12345)
#define HBL_ROOT "ms0:/PSP/SAVEDATA/TODO/"

]

config_keys.each { |key|
	if (config_values[key].size > 0)
		out.puts('#define ' + key + ' { ' + config_values[key].join(', ')  + ' }')
	end
}
